{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vkdnLiKk71g-"
      },
      "source": [
        "##### Copyright 2024 Google Inc."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "0asMuNro71hA"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZOoO2k9AS_ix"
      },
      "source": [
        "# FAX - Federated primitives in JAX\n",
        "\n",
        "In this colab, we will learn what FAX is, how to use it, and why it was designed. We'll go over a couple of JAX-related topics on the way, but some minor familiarity with JAX may be useful."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wArQh8SwVKu1"
      },
      "source": [
        "# Imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "091t-CEOxmNU"
      },
      "outputs": [],
      "source": [
        "!pip install --upgrade google-fax"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KRawRHQeJ1hW"
      },
      "outputs": [],
      "source": [
        "import jax\n",
        "import jax.numpy as jnp\n",
        "import fax"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Xrwn2ZHyTTt6"
      },
      "source": [
        "# What is FAX?\n",
        "\n",
        "FAX has a few goals.\n",
        "\n",
        "1. Create a JAX authoring surface for FL research that uses TFF-like primitives.\n",
        "2. Enable the use of [Federated Automatic Differentiation](https://arxiv.org/abs/2301.07806) (federated AD)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ow6te-MvUPTS"
      },
      "source": [
        "## What is federated AD?\n",
        "\n",
        "Suppose I have (in the parlance of TFF), a **federated computation** with the following type signature:\n",
        "```\n",
        "foo: (x@SERVER, y@CLIENTS) -\u003e float32@SERVER\n",
        "```\n",
        "This looks like a server's loss function - it takes in some server value, and outputs a server float. In many settings, we want to compute the **derivative** of this function - `dz/dx`. This allows us to do things like gradient descent.\n",
        "\n",
        "For non-federated computations, in frameworks like TF, PyTorch, JAX, we can just call `grad(foo)` For example, in JAX we can do the following:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IqI_D5bvUoUh"
      },
      "outputs": [],
      "source": [
        "def square_and_dot(x, y):\n",
        "  return jnp.dot(jnp.square(x), y)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "executionInfo": {
          "elapsed": 56,
          "status": "ok",
          "timestamp": 1706149717284,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": 480
        },
        "id": "MKWJKlTdVzsg",
        "outputId": "965e9495-709f-4b6c-84cd-0c5882949b25"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array(20., dtype=float32)"
            ]
          },
          "execution_count": 95,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "x = jnp.array([1.0, -3.0])\n",
        "y = jnp.array([2.0, 2.0])\n",
        "square_and_dot(x, y)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UaLKVVxtVBJb"
      },
      "source": [
        "To get the derivative of the output with respect to `x`, we can just use `jax.grad`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "executionInfo": {
          "elapsed": 55,
          "status": "ok",
          "timestamp": 1706149717591,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": 480
        },
        "id": "scFCqL9xU7G5",
        "outputId": "025e3ef9-2de2-40fb-f86d-624170d08726"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array([  4., -12.], dtype=float32)"
            ]
          },
          "execution_count": 96,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jax.grad(square_and_dot)(x, y)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FEfs4KOFWAUP"
      },
      "source": [
        "With a single call to `jax.grad`, we can compute the derivative. This is using what's known as **automatic differentiation** (AD).\n",
        "\n",
        "We would like to be able to use this for **federated computation**. While the [federated AD paper](https://arxiv.org/abs/2301.07806) gives a theoretical framework for doing this, it does not have any direct implementation.\n",
        "\n",
        "This is where FAX comes in. FAX defines federated primitives (eg. `federated_broadcast`, `federated_map`, etc.) as JAX primitives. This allows JAX to differentiate through them automatically, enabling federated AD!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eDv6CtR2WnLN"
      },
      "source": [
        "# FAX and Federated Primitives"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bcXrZIDLWsa2"
      },
      "source": [
        "Let's take a look at how we can define federated computations in FAX. Recall that in TFF we have 2 placements, `SERVER` and `CLIENTS`. These are represented as a singleton and as a list, respectively.\n",
        "\n",
        "FAX takes a similar approach - `SERVER` placed values are a singleton."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "executionInfo": {
          "elapsed": 3,
          "status": "ok",
          "timestamp": 1706149717845,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": 480
        },
        "id": "hfVwpWCDXNwa",
        "outputId": "da65afe9-c24a-4edd-d492-6936a4154662"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(2,)"
            ]
          },
          "execution_count": 97,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "server_value = jnp.array([-1.0, 1.0])\n",
        "server_value.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "06S7lGrWXODX"
      },
      "source": [
        "For `CLIENTS` placements, instead of using a list, we add an extra axis to our tensors. This axis represents our clients."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "executionInfo": {
          "elapsed": 3,
          "status": "ok",
          "timestamp": 1706149718145,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": 480
        },
        "id": "FbF8WdB6W98X",
        "outputId": "4a324640-fb86-47ec-e666-bd318194c988"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(3, 2)"
            ]
          },
          "execution_count": 98,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "client_values = jnp.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])\n",
        "client_values.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I41HBhtzXDj3"
      },
      "source": [
        "In the above code, we have an example with 3 clients, each of which has a 1-d tensor of shape `(2,)`. The server also has a 1-d tensor of shape `(2,)`. Let's define a map to transform such vectors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "executionInfo": {
          "elapsed": 4,
          "status": "ok",
          "timestamp": 1706149718439,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": 480
        },
        "id": "0sf7VOI7XX_H",
        "outputId": "29bfd432-f1cf-430b-eda8-7caad70b8617"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "[0. 3.]\n",
            "[[2. 4.]\n",
            " [4. 6.]\n",
            " [6. 8.]]\n"
          ]
        }
      ],
      "source": [
        "def add_constant(x):\n",
        "  return x + jnp.array([1.0, 2.0])\n",
        "\n",
        "print(add_constant(server_value))\n",
        "print(add_constant(client_values))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-4JUrbiAXkrp"
      },
      "source": [
        "We see something interesting here - JAX can recognize additional axes, as in the case of the CLIENTS value, and automatically tries to vectorize."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Wh8zvpK-X1mx"
      },
      "source": [
        "Now let's try to apply a `federated_map`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "executionInfo": {
          "elapsed": 53,
          "status": "ok",
          "timestamp": 1706149718785,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": 480
        },
        "id": "gVA-6c7dXZ3L",
        "outputId": "2f57707c-3ce0-4c90-f2a1-5bbe23def38e"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array([[2., 4.],\n",
              "       [4., 6.],\n",
              "       [6., 8.]], dtype=float32)"
            ]
          },
          "execution_count": 100,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "@fax.fax_program(placements={'clients': 3})\n",
        "def clients_add_one(x):\n",
        "  return fax.federated_map_clients(add_constant, x)\n",
        "\n",
        "clients_add_one(client_values)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "30tVSSZSZXba"
      },
      "source": [
        "We see the same thing! Let's try another primitive, `federated_broadcast`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "executionInfo": {
          "elapsed": 54,
          "status": "ok",
          "timestamp": 1706149719104,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": 480
        },
        "id": "8KgTNTXPX9wr",
        "outputId": "3486691b-2458-483e-b033-fa722a6cce30"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array([[-1.,  1.],\n",
              "       [-1.,  1.],\n",
              "       [-1.,  1.]], dtype=float32)"
            ]
          },
          "execution_count": 101,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "@fax.fax_program(placements={'clients': 3})\n",
        "def broadcast(x):\n",
        "  return fax.federated_broadcast(x)\n",
        "\n",
        "broadcast(server_value)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SVZZdc-xZgcm"
      },
      "source": [
        "This does what we expect - we send the same vector to all clients. So far, no real surprises. The cool thing though, is that we can apply differentiation, just as we did above.\n",
        "\n",
        "The only real restriction is that we want to differentiate server values, with respect to server values. So let's do that."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "executionInfo": {
          "elapsed": 76,
          "status": "ok",
          "timestamp": 1706149719434,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": 480
        },
        "id": "vEkNGpTYZuJK",
        "outputId": "88bf9156-982b-4e4f-e7f6-55c232f0be35"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array(6., dtype=float32)"
            ]
          },
          "execution_count": 102,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "@fax.fax_program(placements={'clients': 3})\n",
        "def broadcast_and_sum(x):\n",
        "  client_x = fax.federated_broadcast(x)\n",
        "  return fax.federated_sum(client_x)\n",
        "\n",
        "broadcast_and_sum(2.0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yyPnPtw5Z8pN"
      },
      "source": [
        "Now we differentiate using *forward-mode differentiation* (more on that at the end). We have to tell JAX which arg to differentiate with respect to, for posterity."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "executionInfo": {
          "elapsed": 103,
          "status": "ok",
          "timestamp": 1706149719795,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": 480
        },
        "id": "uBNPvVj8Z5u5",
        "outputId": "c4359a9f-8e9d-4e9c-ba42-35598955b5e2"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array(3., dtype=float32)"
            ]
          },
          "execution_count": 103,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jax.grad(broadcast_and_sum, argnums=0)(2.0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6JrQoRBbZkIH"
      },
      "source": [
        "Why is this the derivative? Well, let's think about what this function is. We get something like:\n",
        "\n",
        "```\n",
        "x -\u003e [x, x, x] -\u003e sum([x, x, x]) = 3x\n",
        "```\n",
        "Taking a derivative with respect to `x`, we should get 3!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wgUfHsFDWfUf"
      },
      "source": [
        "# Federated Linear Regression and Federated AD"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "InndTb0EWiUV"
      },
      "source": [
        "Let's see how this all works in a more interesting example. We're going to do something akin to linear regression.\n",
        "\n",
        "Let's assume all clients have their own 2d vector `y`. Given a 2d linear regression model `x`, we'll set up our objective function as follows:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rlHgd2_fL2P3"
      },
      "outputs": [],
      "source": [
        "def compute_loss(x, y):\n",
        "  return 0.5*jnp.square(jnp.dot(x, y) - 1.0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hKwDxb6fbCaH"
      },
      "source": [
        "Essentially, this is doing linear regression where (1) each client has a single example and (2) all clients have label 1.0 for that example. This is not an important observation, the point is that `x, y` go in, and we get out some scalar loss.\n",
        "\n",
        "Let's try it out."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "executionInfo": {
          "elapsed": 73,
          "status": "ok",
          "timestamp": 1706149720434,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": 480
        },
        "id": "XNjSjJ-ePwk8",
        "outputId": "72eda4dd-beef-4a94-bc48-308c88436aec"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array(364.5, dtype=float32)"
            ]
          },
          "execution_count": 105,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "compute_loss(jnp.array([2.0, 3.0]), jnp.array([5.0, 6.0]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uY6cJzBVbScR"
      },
      "source": [
        "Great! Now, we can do a `federated_eval` with this loss function. To do that we will:\n",
        "\n",
        "1. Broadcast `x` to the clients.\n",
        "2. Do `compute_loss(x, y)` for each client.\n",
        "3. Average the results.\n",
        "\n",
        "We can do that in FAX as follows."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "executionInfo": {
          "elapsed": 185,
          "status": "ok",
          "timestamp": 1706149720871,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": 480
        },
        "id": "b8_s9bh8PydP",
        "outputId": "bda2597c-e5f6-4132-b995-98af82a5161e"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array(87.16667, dtype=float32)"
            ]
          },
          "execution_count": 106,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "@fax.fax_program(placements={'clients': 3})\n",
        "def federated_eval(server_vector, client_vectors):\n",
        "  broadcast_vector = fax.federated_broadcast(server_vector)\n",
        "  client_losses = fax.federated_map_clients(compute_loss, (broadcast_vector, client_vectors))\n",
        "  return fax.federated_mean(client_losses)\n",
        "\n",
        "server_vector = jnp.array([2.0, -1.0])\n",
        "client_vectors = jnp.array([[1.0, 2.0], [3.0, -4.0], [-7.0, 6.0]])\n",
        "federated_eval(server_vector, client_vectors)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0c6I7j2QbkXS"
      },
      "source": [
        "Just as above, we can use `jax.grad` to differentiate through this function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "executionInfo": {
          "elapsed": 292,
          "status": "ok",
          "timestamp": 1706149721470,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": 480
        },
        "id": "0SRZMRPuR6mu",
        "outputId": "03f59eb7-96f1-478d-e65b-76039eabb271"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array([ 57.666668, -54.666668], dtype=float32)"
            ]
          },
          "execution_count": 107,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "grad_fn = jax.grad(federated_eval, argnums=0)\n",
        "grad_fn(server_vector, client_vectors)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D9HcGfN8bofE"
      },
      "source": [
        "Let's pause here to think about what this means.\n",
        "\n",
        "We have some loss function $\\ell(x, y)$. Our federated evaluation is computing, for some set of clients $C$:\n",
        "\n",
        "$$\n",
        "\\dfrac{1}{|C|}\\sum_{i \\in C} \\ell(x, y_i).\n",
        "$$\n",
        "\n",
        "By taking a derivative we get:\n",
        "\n",
        "$$\n",
        "\\dfrac{d}{dx}\\left(\\dfrac{1}{|C|}\\sum_{i \\in C} \\ell(x, y_i) \\right) = \\dfrac{1}{|C|} \\sum_{i \\in C} \\dfrac{d\\ell}{dx}(x, y_i).\n",
        "$$\n",
        "\n",
        "In other words we are just computing **the average gradient across clients**. This is exactly what FedSGD does! The only missing ingredient is to use that derivative to update our model. We'll do that via gradient descent with learning rate 0.01."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "executionInfo": {
          "elapsed": 206,
          "status": "ok",
          "timestamp": 1706149721928,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": 480
        },
        "id": "znvCCrWYSeu5",
        "outputId": "eeb1cefc-2c0e-4f91-b8b9-d0e351024e71"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "87.16667\n",
            "35.551258\n"
          ]
        }
      ],
      "source": [
        "@fax.fax_program(placements={'clients': 3})\n",
        "def fed_sgd_step(server_vector, client_vectors):\n",
        "  server_grad = grad_fn(server_vector, client_vectors)\n",
        "  return server_vector - 0.01 * server_grad\n",
        "\n",
        "print(federated_eval(server_vector, client_vectors))\n",
        "updated_server_vector = fed_sgd_step(server_vector, client_vectors)\n",
        "print(federated_eval(updated_server_vector, client_vectors))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bSX-tgF8dFX4"
      },
      "source": [
        "As expected, after applying a step of FedSGD, we get a vector with lower loss!\n",
        "\n",
        "Moreover, going from `federated_eval -\u003e FedSGD` was essentially trivial - we just applied federated AD!"
      ]
    },
    {
      "metadata": {
        "id": "7Sle5BNYSRTi"
      },
      "cell_type": "markdown",
      "source": [
        "# Conclusion\n",
        "\n",
        "Above, we showed how to use FAX to define federated computations, and how to apply federated AD to differentiate through them. We encourage you to try out your own federated computations, especially at scale."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
